// lexer.c
#include "Yecc.h"

const struct KeywordInfo keyword_table[] = {
    {"void",     KW_VOID},
    {"bool",     KW_BOOL},
    {"char",     KW_CHAR},
    {"short",    KW_SHORT},
    {"int",      KW_INT},
    {"long",     KW_LONG},
    {"enum",     KW_ENUM},
    {"struct",   KW_STRUCT},
    {"const",    KW_CONST},
    {"signed",   KW_SIGNED},
    {"unsigned", KW_UNSIGNED},
    {"sizeof",   KW_SIZEOF},
    {"if",       KW_IF},
    {"else",     KW_ELSE},
    {"switch",   KW_SWITCH},
    {"case",     KW_CASE},
    {"default",  KW_DEFAULT},
    {"do",       KW_DO},
    {"for",      KW_FOR},
    {"while",    KW_WHILE},
    {"break",    KW_BREAK},
    {"continue", KW_CONTINUE},
    {"return",   KW_RETURN},
    {"asm",      KW_ASM},
    {"true",     KW_TRUE},
    {"false",    KW_FALSE},
    {"NULL",     KW_NULL},
    {NULL,       KW_UNKNOWN}
};

// for longest match
const struct PunctInfo punct_table[] = {
    // 3-char
    {"<<=", PK_SHL_ASG},
    {">>=", PK_SHR_ASG},
    {"...", PK_ELLIPSIS},
    // 2-char
    {"&=",  PK_AND_ASG}, 
    {"&&",  PK_LOG_AND},
    {"|=",  PK_OR_ASG},  
    {"||",  PK_LOG_OR},
    {"^=",  PK_XOR_ASG},
    {"<=",  PK_LE}, 
    {"<<",  PK_SHL}, 
    {">=",  PK_GE}, 
    {">>",  PK_SHR},
    {"==",  PK_EQ}, 
    {"!=",  PK_NEQ},
    {"+=",  PK_ADD_ASG}, 
    {"++",  PK_INC},
    {"-=",  PK_SUB_ASG}, 
    {"--",  PK_DEC}, 
    {"->",  PK_ARROW},
    {"*=",  PK_MUL_ASG}, 
    {"/=",  PK_DIV_ASG}, 
    {"%=",  PK_MOD_ASG},
    // 1-char
    {"(",   PK_LPAREN}, 
    {")",   PK_RPAREN},
    {"[",   PK_LBRACK}, 
    {"]",   PK_RBRACK},
    {"{",   PK_LBRACE}, 
    {"}",   PK_RBRACE},
    {",",   PK_COMMA}, 
    {":",   PK_COLON}, 
    {";",   PK_SEMI}, 
    {"?",   PK_QUESTION},
    {"~",   PK_BITNOT},
    {"+",   PK_ADD}, 
    {"-",   PK_SUB}, 
    {"*",   PK_MUL}, 
    {"/",   PK_DIV}, 
    {"%",   PK_MOD},
    {"&",   PK_BITAND}, 
    {"|",   PK_BITOR}, 
    {"^",   PK_BITXOR},
    {"<",   PK_LT}, 
    {">",   PK_GT}, 
    {"=",   PK_ASSIGN}, 
    {"!",   PK_LOG_NOT}, 
    {".",   PK_DOT},
    {NULL,  PK_UNKNOWN}
};

int Len(const char *str) {
    int cnt = 0;
    while (str[cnt] != '\0') {
        cnt++;
    }
    return cnt;
}

void Cpy(char *dst, const char *src) {
    int i = 0;
    while (src[i] != '\0') {
        dst[i] = src[i];
        i++;
    }
    dst[i] = '\0';
    return ;
}

void Cat(char *dst, const char *src) {
    int i = 0;
    while(dst[i] != '\0') i++;
    int j = 0;
    while(src[j] != '\0') {
        dst[i + j] = src[j];
        j++;
    }
    dst[i + j] = '\0';
    return;
}

int Cmp(const char *s1, const char *s2, const int len) {
    int i = 0;
    while (i < len) {
        bool dif = s1[i] != s2[i];
        bool s1e = s1[i] == '\0';
        bool s2e = s2[i] == '\0';
        if (dif || s1e || s2e) return s1[i] - s2[i];
        i++;
        
    }
    return 0;
}

unsigned long Atoi(const char* str, int base) {
    int i = 0;
    unsigned long res = 0;
    if (!str[0]) {
        printf("empty input for atoi\n");
        assert(0);
    }
    while (str[i] != '\0') {
        unsigned char digit = 0;
        if (str[i] >= '0' && str[i] <= '9') digit = str[i] - '0';
        else if (str[i] >= 'A' && str[i] <= 'F') digit = str[i] - 'A' + 10;
        else if (str[i] >= 'a' && str[i] <= 'f') digit = str[i] - 'a' + 10;
        else {
            printf("illegal char '%c' in atoi\n", str[i]);
            assert(0);
        }
        if (digit >= base) {
            printf("digit '%c' out of base %d in atoi\n", str[i], base);
            assert(0);
        }
        res = res * base + digit;
        i++;
    }
    return res;
}

void Itoa(char* buf, long data) {

    assert(buf != NULL);
    char tmp[32];
    int i = 0;
    int sign = 0;

    // for sign
    if(data < 0) {
        sign = 1;
        data = -data;
    }

    // for 0
    if(data == 0) {
        buf[0] = '0';
        buf[1] = '\0';
        return;
    }

    // extract data (reverse)
    while(data > 0) {
        tmp[i++] = '0' + (data % 10);
        data /= 10;
    }

    // add sign
    int pos = 0;
    if(sign) buf[pos++] = '-';
    // reverse and output to buf
    while(i > 0)
        buf[pos++] = tmp[--i];

    buf[pos] = '\0';

    return ;
}

void dumpText(char *buf, const char* src, const int loc, const int len) {
    int i = 0;
    while(i < len && i < 255) {
        buf[i] = src[loc + i];
        if(buf[i] == '\0') return ;
        i++;
    }
    buf[i] = '\0';
    return ;
}

enum PunctKind getPunctKind(const char* str) {
    enum PunctKind pk = PK_UNKNOWN;
    for (int i = 0; punct_table[i].text != NULL; i++) {
        int len = Len(punct_table[i].text);
        // punct len does not match
        if(len != Len(str)) continue;
        if (Cmp(str, punct_table[i].text, len) == 0) {
            return punct_table[i].pk;
        }
    }
    return pk;
}

enum KeywordKind getKeywordKind(const char* str) {
    enum KeywordKind kw = KW_UNKNOWN;
    for (int i = 0; keyword_table[i].text != NULL; i++) {
        int len = Len(keyword_table[i].text);
        if(len != Len(str)) continue;
        if (Cmp(keyword_table[i].text, str, len) == 0)
            return keyword_table[i].kw;
    }
    return kw;
}

enum PunctKind isPunct(const char* src, int loc) {
    enum PunctKind pk = PK_UNKNOWN;
    char buf[8];
    for (int i = 0; punct_table[i].text != NULL; i++) {
        int len = Len(punct_table[i].text);
        dumpText(buf, src, loc, len);
        if (getPunctKind(buf) == punct_table[i].pk) return punct_table[i].pk;
    }
    return pk;
}

int getPunctLength(enum PunctKind pk) {
    for (int i = 0; punct_table[i].text != NULL; i++) {
        if (punct_table[i].pk == pk)
            return Len(punct_table[i].text);
    }
    return 0;
}

const char* getPunctText(enum PunctKind pk) {
    for (int i = 0; punct_table[i].text != NULL; i++) {
        if (punct_table[i].pk == pk)
            return punct_table[i].text;
    }
    return NULL;
}

// only \n \0 \\ \' \"
char isEscapeChar(const char* src, int idx) {
    assert(src[idx] == '\\');
    idx++; // skip the '\\'
    char ch = 0;
    switch (src[idx]) {
        case 'n':  ch = '\n'; break;
        case '0':  ch = '\0'; break;
        case '\'': ch = '\''; break;
        case '"':  ch = '"';  break;
        case '\\': ch = '\\'; break;
        default:
            printf("Error: unsupported escape sequence '\\%c' \n", src[idx]);
            assert(0);
    }
    return ch;
}

bool isDigit(const char ch) {
    return ch >= '0' && ch <= '9';
}

bool isBinDigit(const char ch) {
    return ch == '0' || ch == '1';
}

bool isHexDigit(const char ch) {
    return isDigit(ch) || (ch >= 'A' && ch <= 'F');
}

bool isLetter(const char ch) {
    bool a = (ch >= 'a' && ch <= 'z');
    bool A = (ch >= 'A' && ch <= 'Z');
    return a || A || ch == '_';
}

bool isSkipChar(const char ch) {
    return ch == ' ' || ch == '\n';
}

struct Token* new_token(enum TokenKind kind) {
    struct Token* tok = (struct Token*)malloc(sizeof(struct Token));
    assert(tok != NULL);
    tok->kind = kind;
    tok->kw = KW_UNKNOWN;
    tok->pk = PK_UNKNOWN;
    tok->num = 0;
    tok->ch = '\0';
    tok->str = NULL;
    tok->name = NULL;
    tok->next = NULL;
    return tok;
}

void check_keyword(struct Token* head) {
    struct Token* tok = head;
    while(tok != NULL) {
        if(tok->kind == TK_IDENT) {
            for(int i = 0; keyword_table[i].text != NULL; i++) {
                int len_kw = Len(keyword_table[i].text);
                int len_id = Len(tok->name);
                if(len_kw != len_id) continue;
                if(Cmp(tok->name, keyword_table[i].text, len_kw) == 0) {
                    tok->kind = TK_KEYWORD;
                    tok->kw = keyword_table[i].kw;
                    break;
                }
            }
        }
        tok = tok->next;
    }
    return ;
}

struct Token* Lexer(const char* src) {

    debug_info_dumpSrc(src);

    struct Token* cur = new_token(TK_HEAD);
    struct Token* head = cur;

    int idx = 0;
    int loc = 0;

    while (src[idx] != '\0') {

        // skip the space or \n
        if(isSkipChar(src[idx])) {
            idx++;
            continue;
        }
        
        // =======================================
        // TK_CH
        // =======================================
        if(src[idx] == '\'') {

            idx++; // skip the opening '\''
            char ch = 0;
            
            if(src[idx] == '\\') {
                ch = isEscapeChar(src, idx);
                idx+=2; // skip the '\\' and the characters
            } else {
                // regular characters
                ch = src[idx];
                idx++; // skip the characters
            }

            assert(src[idx] == '\'');
            idx++; // skip the closing '\''

            // create new token
            cur->next = new_token(TK_CH);
            cur = cur->next;
            // fill the critical field: ch
            cur->ch = ch;
            continue;
        }

        // =======================================
        // TK_STR
        // =======================================
        if (src[idx] == '"') {

            idx++; // skip the opening '"'
            char buf[4096];
            int bpos = 0;

            while (src[idx] != '"') {
                assert(src[idx] != '\0');
                if (src[idx] == '\\') {
                    buf[bpos++] = isEscapeChar(src, idx);
                    idx += 2; // skip '\' + actual escaped char
                } else {
                    buf[bpos++] = src[idx++];
                }
                assert(bpos <= 4094);
            }

            assert(src[idx] == '"');
            idx++; // skip the closing '"'
            buf[bpos] = '\0';
            
            // create new token
            cur->next = new_token(TK_STR);
            cur = cur->next;
            // fill the critical field: str
            cur->str = malloc(bpos * sizeof(char));
            assert(cur->str != NULL);
            Cpy(cur->str, buf);
            assert(Len(cur->str) == bpos);
            continue;
        }

        // =======================================
        // TK_NUM
        // =======================================
        if(isDigit(src[idx])) {
            loc = idx;
            int base = 10;
            if(src[idx] == '0') {
                idx++;
                if(src[idx] == 'b') {
                    // bin
                    idx++;  // skip "b"
                    loc+=2; // skip "0b"
                    base = 2;
                    while(isBinDigit(src[idx])) idx++;
                } else if(src[idx] == 'x') {
                    // hex
                    idx++;  // skip "x"
                    loc+=2; // skip "0x"
                    base = 16;
                    while(isHexDigit(src[idx])) idx++;
                } else {
                    // just num 0, do nothing
                }
            } else {
                // dec
                while(isDigit(src[idx])) idx++;
            }
            
            // create new token
            cur->next = new_token(TK_NUM);
            cur = cur->next;
            // fill the critical field: num
            char buf[256];
            int num_len = idx - loc;
            dumpText(buf, src, loc, num_len);
            cur->num = Atoi(buf, base);
            continue;
        }

        // =======================================
        // TK_IDENT
        // =======================================
        if(isLetter(src[idx])) { // isLetter include '_'
            loc = idx;
            while(isLetter(src[idx]) || isDigit(src[idx])) idx++;
            
            // create new token
            cur->next = new_token(TK_IDENT);
            cur = cur->next;
            // fill the critical field: str
            int ident_len = idx - loc;
            cur->name = malloc(ident_len * sizeof(char));
            assert(cur->name != NULL);
            dumpText(cur->name, src, loc, ident_len);
            assert(Len(cur->name) == ident_len);
            continue;
        }

        // =======================================
        // TK_PUNCT
        // =======================================

        enum PunctKind pk = isPunct(src, idx);
        if(pk != PK_UNKNOWN) {
            int punct_len = getPunctLength(pk);
            assert(punct_len != 0);
            idx += punct_len;
            // create new token
            cur->next = new_token(TK_PUNCT);
            cur = cur->next;
            // fill the critical field: pk
            cur->pk = pk;
            continue;
        }

        printf("Unknown token: ");
        for(int i = 0; i < 100; i++) {
            printf("%c", src[idx++]);
        }
        assert(0);
    }
    
    check_keyword(head);

    cur->next = new_token(TK_EOF);

    printf("### Lexer Completed! ###\n\n");

    debug_info_dumpTokens(head);
    
    return head;
}
